#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
会话管理器
管理用户会话和任务的关系，支持前端刷新后取消任务
"""

import time
import logging
import threading
from typing import Dict, Set, Optional, List
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

class SessionManager:
    """会话管理器"""
    
    def __init__(self, session_timeout_minutes: int = 30):
        """
        初始化会话管理器
        
        Args:
            session_timeout_minutes: 会话超时时间（分钟）
        """
        self.session_timeout_minutes = session_timeout_minutes
        self.sessions: Dict[str, Dict] = {}  # session_id -> session_info
        self.session_tasks: Dict[str, Set[str]] = {}  # session_id -> task_ids
        self.task_sessions: Dict[str, str] = {}  # task_id -> session_id
        self.lock = threading.Lock()
        
        # 启动清理线程
        self.cleanup_thread = threading.Thread(target=self._cleanup_expired_sessions, daemon=True)
        self.cleanup_thread.start()
        
        logger.info(f"会话管理器已初始化，超时时间: {session_timeout_minutes}分钟")
    
    def create_session(self, session_id: str, client_ip: str) -> bool:
        """
        创建新会话
        
        Args:
            session_id: 会话ID
            client_ip: 客户端IP
            
        Returns:
            是否创建成功
        """
        with self.lock:
            if session_id in self.sessions:
                # 更新现有会话
                self.sessions[session_id]['last_activity'] = time.time()
                self.sessions[session_id]['client_ip'] = client_ip
                logger.info(f"会话已更新: {session_id}")
                return True
            
            # 创建新会话
            self.sessions[session_id] = {
                'created_at': time.time(),
                'last_activity': time.time(),
                'client_ip': client_ip,
                'task_count': 0
            }
            self.session_tasks[session_id] = set()
            
            logger.info(f"新会话已创建: {session_id} (IP: {client_ip})")
            return True
    
    def add_task_to_session(self, session_id: str, task_id: str) -> bool:
        """
        将任务添加到会话
        
        Args:
            session_id: 会话ID
            task_id: 任务ID
            
        Returns:
            是否添加成功
        """
        with self.lock:
            if session_id not in self.sessions:
                logger.warning(f"会话不存在: {session_id}")
                return False
            
            # 更新会话活动时间
            self.sessions[session_id]['last_activity'] = time.time()
            self.sessions[session_id]['task_count'] += 1
            
            # 添加任务到会话
            self.session_tasks[session_id].add(task_id)
            self.task_sessions[task_id] = session_id
            
            logger.info(f"任务已添加到会话: {task_id} -> {session_id}")
            return True
    
    def remove_task_from_session(self, task_id: str) -> bool:
        """
        从会话中移除任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            是否移除成功
        """
        with self.lock:
            if task_id not in self.task_sessions:
                return False
            
            session_id = self.task_sessions[task_id]
            
            # 从会话中移除任务
            if session_id in self.session_tasks:
                self.session_tasks[session_id].discard(task_id)
            
            # 更新会话统计
            if session_id in self.sessions:
                self.sessions[session_id]['task_count'] = max(0, 
                    self.sessions[session_id]['task_count'] - 1)
            
            # 移除任务-会话映射
            del self.task_sessions[task_id]
            
            logger.info(f"任务已从会话移除: {task_id} -> {session_id}")
            return True
    
    def get_session_tasks(self, session_id: str) -> Set[str]:
        """
        获取会话的所有任务
        
        Args:
            session_id: 会话ID
            
        Returns:
            任务ID集合
        """
        with self.lock:
            return self.session_tasks.get(session_id, set()).copy()
    
    def get_task_session(self, task_id: str) -> Optional[str]:
        """
        获取任务所属的会话
        
        Args:
            task_id: 任务ID
            
        Returns:
            会话ID，如果不存在则返回None
        """
        with self.lock:
            return self.task_sessions.get(task_id)
    
    def is_session_active(self, session_id: str) -> bool:
        """
        检查会话是否活跃
        
        Args:
            session_id: 会话ID
            
        Returns:
            是否活跃
        """
        with self.lock:
            if session_id not in self.sessions:
                return False
            
            last_activity = self.sessions[session_id]['last_activity']
            timeout_seconds = self.session_timeout_minutes * 60
            
            return (time.time() - last_activity) < timeout_seconds
    
    def cancel_session_tasks(self, session_id: str) -> List[str]:
        """
        取消会话的所有任务
        
        Args:
            session_id: 会话ID
            
        Returns:
            被取消的任务ID列表
        """
        with self.lock:
            if session_id not in self.session_tasks:
                return []
            
            tasks_to_cancel = list(self.session_tasks[session_id])
            
            # 清空会话任务
            self.session_tasks[session_id].clear()
            
            # 移除任务-会话映射
            for task_id in tasks_to_cancel:
                if task_id in self.task_sessions:
                    del self.task_sessions[task_id]
            
            # 更新会话统计
            if session_id in self.sessions:
                self.sessions[session_id]['task_count'] = 0
            
            logger.info(f"会话任务已取消: {session_id} ({len(tasks_to_cancel)}个任务)")
            return tasks_to_cancel
    
    def _cleanup_expired_sessions(self):
        """清理过期会话的后台线程"""
        while True:
            try:
                time.sleep(60)  # 每分钟检查一次
                
                current_time = time.time()
                timeout_seconds = self.session_timeout_minutes * 60
                
                with self.lock:
                    expired_sessions = []
                    for session_id, session_info in self.sessions.items():
                        if (current_time - session_info['last_activity']) > timeout_seconds:
                            expired_sessions.append(session_id)
                    
                    for session_id in expired_sessions:
                        # 取消过期会话的任务
                        tasks_to_cancel = self.cancel_session_tasks(session_id)
                        
                        # 移除过期会话
                        del self.sessions[session_id]
                        if session_id in self.session_tasks:
                            del self.session_tasks[session_id]
                        
                        logger.info(f"过期会话已清理: {session_id} ({len(tasks_to_cancel)}个任务)")
                
            except Exception as e:
                logger.error(f"清理过期会话时出错: {e}")
    
    def get_session_stats(self) -> Dict:
        """
        获取会话统计信息
        
        Returns:
            统计信息字典
        """
        with self.lock:
            active_sessions = sum(1 for session_info in self.sessions.values() 
                                if self.is_session_active(session_info.get('session_id', '')))
            
            total_tasks = sum(len(tasks) for tasks in self.session_tasks.values())
            
            return {
                'total_sessions': len(self.sessions),
                'active_sessions': active_sessions,
                'total_tasks_in_sessions': total_tasks,
                'session_timeout_minutes': self.session_timeout_minutes
            }

# 全局实例
session_manager = SessionManager()
